import sqlite3
import os
from datetime import datetime
from passlib.context import CryptContext

# 密码加密上下文（与 main.py 共用）
pwd_context = CryptContext(schemes=["bcrypt"], deprecated="auto")

# 数据库路径（与 main.py 一致）
current_file_path = os.path.abspath(__file__)
app_dir = os.path.dirname(current_file_path)
DATABASE_FILE = os.path.join(app_dir, "data", "deepseek_chat.db")

def init_db():
    """初始化数据库：创建用户表、会话表、聊天消息表、API Key表"""
    # 创建数据目录
    storage_dir = os.path.dirname(DATABASE_FILE)
    if not os.path.exists(storage_dir):
        os.makedirs(storage_dir)
    
    conn = sqlite3.connect(DATABASE_FILE)
    cursor = conn.cursor()
    
    # 1. 用户表（新增 avatar 字段存储头像Base64）
    cursor.execute('''
    CREATE TABLE IF NOT EXISTS users (
        id INTEGER PRIMARY KEY AUTOINCREMENT,
        username TEXT UNIQUE NOT NULL,
        password_hash TEXT NOT NULL,
        email TEXT UNIQUE,
        is_admin BOOLEAN DEFAULT 0,
        avatar TEXT,  -- 存储头像Base64字符串
        created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
    )
    ''')
    
    # 2. 聊天会话表（用于存储历史话题列表）
    cursor.execute('''
    CREATE TABLE IF NOT EXISTS chat_sessions (
        id INTEGER PRIMARY KEY AUTOINCREMENT,
        user_id INTEGER NOT NULL,
        title TEXT NOT NULL,  -- 会话标题（如“Python代码优化”）
        created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
        updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
        FOREIGN KEY (user_id) REFERENCES users (id) ON DELETE CASCADE
    )
    ''')
    
    # 3. 聊天消息表（关联会话）
    cursor.execute('''
    CREATE TABLE IF NOT EXISTS chat_messages (
        id INTEGER PRIMARY KEY AUTOINCREMENT,
        session_id INTEGER NOT NULL,
        user_id INTEGER NOT NULL,
        role TEXT NOT NULL,  -- 'user' 或 'assistant'
        content TEXT NOT NULL,
        created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
        FOREIGN KEY (session_id) REFERENCES chat_sessions (id) ON DELETE CASCADE,
        FOREIGN KEY (user_id) REFERENCES users (id) ON DELETE CASCADE
    )
    ''')
    
    # 4. API Key表（补充激活管理员ID字段）
    cursor.execute('''
    CREATE TABLE IF NOT EXISTS api_keys (
        id INTEGER PRIMARY KEY AUTOINCREMENT,
        user_id INTEGER NOT NULL,
        key TEXT UNIQUE NOT NULL,
        description TEXT,
        is_active BOOLEAN DEFAULT 0,  -- 0:待审核/1:已激活
        is_revoked BOOLEAN DEFAULT 0, -- 0:正常/1:已注销
        created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
        activated_at TIMESTAMP NULL,
        activated_by INTEGER NULL,  -- 激活管理员ID
        revoked_at TIMESTAMP NULL,
        revoked_by INTEGER NULL,    -- 注销管理员ID
        FOREIGN KEY (user_id) REFERENCES users (id) ON DELETE CASCADE,
        FOREIGN KEY (activated_by) REFERENCES users (id),
        FOREIGN KEY (revoked_by) REFERENCES users (id)
    )
    ''')
    
    # 插入默认管理员（密码：admin123，首次初始化时创建）
    cursor.execute('SELECT * FROM users WHERE username = "admin"')
    if not cursor.fetchone():
        admin_hash = pwd_context.hash("admin123")
        cursor.execute('''
        INSERT INTO users (username, password_hash, email, is_admin)
        VALUES (?, ?, ?, ?)
        ''', ("admin", admin_hash, "admin@example.com", 1))
    
    conn.commit()
    conn.close()

def get_db_connection():
    """获取数据库连接"""
    conn = sqlite3.connect(DATABASE_FILE)
    conn.row_factory = sqlite3.Row  # 支持按列名访问
    return conn

# --------------------------
# 用户相关操作
# --------------------------
def get_user_by_username(username):
    """通过用户名获取用户"""
    conn = get_db_connection()
    user = conn.execute('SELECT * FROM users WHERE username = ?', (username,)).fetchone()
    conn.close()
    return dict(user) if user else None

def get_user_by_id(user_id):
    """通过ID获取用户"""
    conn = get_db_connection()
    user = conn.execute('SELECT * FROM users WHERE id = ?', (user_id,)).fetchone()
    conn.close()
    return dict(user) if user else None

def create_user(username, password_hash, email=None):
    """创建新用户"""
    conn = get_db_connection()
    cursor = conn.cursor()
    cursor.execute('''
    INSERT INTO users (username, password_hash, email)
    VALUES (?, ?, ?)
    ''', (username, password_hash, email))
    user_id = cursor.lastrowid
    conn.commit()
    
    # 获取创建的用户
    user = conn.execute('SELECT * FROM users WHERE id = ?', (user_id,)).fetchone()
    conn.close()
    return dict(user)

def update_user_avatar(user_id, avatar_base64):
    """更新用户头像"""
    conn = get_db_connection()
    conn.execute('UPDATE users SET avatar = ? WHERE id = ?', (avatar_base64, user_id))
    conn.commit()
    conn.close()

def get_all_users():
    """获取所有用户（管理员用）"""
    conn = get_db_connection()
    users = conn.execute('SELECT id, username, email, is_admin, created_at FROM users').fetchall()
    conn.close()
    return [dict(user) for user in users]

def is_admin(user_id):
    """判断用户是否为管理员"""
    user = get_user_by_id(user_id)
    return user and user["is_admin"] == 1

# --------------------------
# 聊天会话相关操作（前端历史话题列表）
# --------------------------
def create_chat_session(user_id, title="新话题"):
    """创建新聊天会话"""
    conn = get_db_connection()
    cursor = conn.cursor()
    cursor.execute('''
    INSERT INTO chat_sessions (user_id, title)
    VALUES (?, ?)
    ''', (user_id, title))
    session_id = cursor.lastrowid
    conn.commit()
    
    session = conn.execute('SELECT * FROM chat_sessions WHERE id = ?', (session_id,)).fetchone()
    conn.close()
    return dict(session)

def get_chat_sessions(user_id, limit=10):
    """获取用户的聊天会话列表（按更新时间倒序）"""
    conn = get_db_connection()
    sessions = conn.execute('''
    SELECT id, title, created_at, updated_at 
    FROM chat_sessions 
    WHERE user_id = ? 
    ORDER BY updated_at DESC 
    LIMIT ?
    ''', (user_id, limit)).fetchall()
    conn.close()
    return [dict(session) for session in sessions]

def update_chat_session_title(session_id, new_title):
    """更新会话标题"""
    conn = get_db_connection()
    conn.execute('''
    UPDATE chat_sessions 
    SET title = ?, updated_at = CURRENT_TIMESTAMP 
    WHERE id = ?
    ''', (new_title, session_id))
    conn.commit()
    conn.close()

def delete_chat_session(session_id):
    """删除会话（级联删除消息）"""
    conn = get_db_connection()
    conn.execute('DELETE FROM chat_sessions WHERE id = ?', (session_id,))
    conn.commit()
    conn.close()

# --------------------------
# 聊天消息相关操作
# --------------------------
def save_chat_message(session_id, user_id, role, content):
    """保存聊天消息"""
    conn = get_db_connection()
    cursor = conn.cursor()
    
    # 保存消息
    cursor.execute('''
    INSERT INTO chat_messages (session_id, user_id, role, content)
    VALUES (?, ?, ?, ?)
    ''', (session_id, user_id, role, content))
    msg_id = cursor.lastrowid
    
    # 更新会话最后更新时间
    cursor.execute('''
    UPDATE chat_sessions 
    SET updated_at = CURRENT_TIMESTAMP 
    WHERE id = ?
    ''', (session_id,))
    
    conn.commit()
    
    # 获取消息详情
    msg = conn.execute('SELECT * FROM chat_messages WHERE id = ?', (msg_id,)).fetchone()
    conn.close()
    return dict(msg)

def get_chat_messages(session_id):
    """获取会话的所有消息"""
    conn = get_db_connection()
    messages = conn.execute('''
    SELECT id, role, content, created_at 
    FROM chat_messages 
    WHERE session_id = ? 
    ORDER BY created_at ASC
    ''', (session_id,)).fetchall()
    conn.close()
    return [dict(msg) for msg in messages]

def clear_chat_messages(session_id):
    """清空会话的消息"""
    conn = get_db_connection()
    conn.execute('DELETE FROM chat_messages WHERE session_id = ?', (session_id,))
    conn.commit()
    conn.close()

# --------------------------
# API Key相关操作
# --------------------------
def create_api_key(user_id, key, description):
    """创建API Key（待审核）"""
    conn = get_db_connection()
    cursor = conn.cursor()
    cursor.execute('''
    INSERT INTO api_keys (user_id, key, description)
    VALUES (?, ?, ?)
    ''', (user_id, key, description))
    key_id = cursor.lastrowid
    conn.commit()
    
    api_key = conn.execute('SELECT * FROM api_keys WHERE id = ?', (key_id,)).fetchone()
    conn.close()
    return dict(api_key)

def get_api_keys_for_user(user_id):
    """获取用户的所有API Key"""
    conn = get_db_connection()
    keys = conn.execute('''
    SELECT id, key, description, is_active, is_revoked, created_at, activated_at
    FROM api_keys 
    WHERE user_id = ? 
    ORDER BY created_at DESC
    ''', (user_id,)).fetchall()
    conn.close()
    return [dict(key) for key in keys]

def get_api_key_by_key(key):
    """通过Key字符串获取API Key信息"""
    conn = get_db_connection()
    api_key = conn.execute('SELECT * FROM api_keys WHERE key = ?', (key,)).fetchone()
    conn.close()
    return dict(api_key) if api_key else None

def activate_api_key(key_id, admin_id):
    """激活API Key（管理员操作）"""
    conn = get_db_connection()
    cursor = conn.cursor()
    cursor.execute('''
    UPDATE api_keys 
    SET is_active = 1, activated_at = CURRENT_TIMESTAMP, activated_by = ? 
    WHERE id = ? AND is_revoked = 0
    ''', (admin_id, key_id))
    conn.commit()
    
    # 检查是否更新成功
    api_key = conn.execute('SELECT * FROM api_keys WHERE id = ?', (key_id,)).fetchone()
    conn.close()
    return dict(api_key) if api_key and api_key["is_active"] == 1 else None

def revoke_api_key(key_id, admin_id=None, user_id=None):
    """注销API Key（管理员/用户自己）"""
    conn = get_db_connection()
    cursor = conn.cursor()
    # 管理员可注销任何Key，用户只能注销自己的Key
    if admin_id:
        condition = 'id = ?'
        params = (key_id, admin_id, admin_id)
    else:
        condition = 'id = ? AND user_id = ?'
        params = (key_id, user_id, user_id)
    
    cursor.execute(f'''
    UPDATE api_keys 
    SET is_revoked = 1, is_active = 0, revoked_at = CURRENT_TIMESTAMP, revoked_by = ? 
    WHERE {condition}
    ''', params)
    rowcount = cursor.rowcount
    conn.commit()
    conn.close()
    return rowcount > 0  # 返回是否注销成功

def get_pending_api_keys():
    """获取待激活的API Key（管理员用）"""
    conn = get_db_connection()
    # 关联用户表，获取用户名
    keys = conn.execute('''
    SELECT ak.id, ak.user_id, u.username, ak.description, ak.created_at
    FROM api_keys ak
    JOIN users u ON ak.user_id = u.id
    WHERE ak.is_active = 0 AND ak.is_revoked = 0
    ORDER BY ak.created_at ASC
    ''', ()).fetchall()
    conn.close()
    return [dict(key) for key in keys]